'''自制的网页爬虫,类似wget,webzip的用法.
  对于相对路径,按照网页本身的代码保存到对应位置.
  对于指定根目录外的资源如图片,则同一用_更名并保存到_outside文件夹

  如果需要过滤,则继承本类并重写_filter函数.
  
  初始页面是一个表,可以对表中的每一项url指定其所在层数,
  层数超过指定上限则会终止
  
  by setycyas @2024-05-07
  '''

import os
from urllib.parse import urlparse, urljoin
import re

from bs4 import BeautifulSoup
import requests

class MyWebGet:

  def __init__(self, root, outputDir, urlList, maxLevel = 3, headers = None, encoding = 'utf-8'):
    """初始化
    参数:
    - root: str, 爬虫地址的根目录,含http(s)://,结尾没有/
    - outputDir: str, 保存到文件夹
    - urlList: list, 爬虫开始页面表,格式为[(地址, 初始层数)],如果
      有些地址不需要爬太多的层,可以把初始层数设定到maxLevel
    - maxLevel: int, 最大层数,当层数大于该数字时不做停止
    - headers: {}, headers
    - encoding: str, 页面编码
    """
    ## 复制初始参数
    self._root = root[:-1] if root.endswith('/') else root
    self._outputDir = outputDir
    self._urlList = urlList
    self._maxLevel = maxLevel
    self._headers = headers if headers else {}
    self._encoding = encoding
      
    ## 其他初始参数
    #self.urlHandler = UrlHandler.UrlHandler(root)
    self._htmlRecord = {} # 已下载的页面,{'地址': 下载时的分析层数}
    self._resource = {} # 需下载的资源文件,{'地址': 文件名完整路径}
    self._session = requests.Session() # 会话
    return
    
  def _urlFilter(self, url):
    """过滤url,返回True的时候url才进入下一层处理.
    必须处理扩展名,因为所有url都要分析,此外会过滤一些其他情况.
    重写本函数可以做更多过滤
    """
    # 允许的扩展名
    expand = ['htm', 'html', 'shtml']
    url = url.lower()
    if '#' in url:
      return False
    if (not url.startswith(self._root)):
      return False
    for x in expand:
      if re.search(f'\.{x}\??[0-9_A-Za-z]*$', url):
        return True
    return False
    
  def _resourceFilter(self, url):
    """过滤资源,不返回True的不下载"""
    # 允许的扩展名
    expand = ['jpg', 'png', 'gif', 'bmp', 'css']
    url = url.lower()
    if '#' in url:
      return False
    for x in expand:
      if re.search(f'\.{x}\??[0-9_A-Za-z]*$', url):
        return True
    return False
   
  def _urlParse(self, url, parentUrl):
    """对一个url进行分析,返回:
    {'inRoot': 是否在网站根目录中, 'fullUrl': 完整url, 
    'downPath': 下载到硬盘中的完整路径, 'filename': 下载到硬盘的文件}
    而parentUrl则是其父url
    注意如果在根目录下,'downPath'就是\
    """
    result = {}
    pParse = urlparse(parentUrl) # parentUrl分析
    pFn = re.search(r'([^/]*)$', parentUrl).group(0) # parent文件名
    ## 先还原完整的url
    if url.startswith('http://') or url.startswith('https://'):
      # 完整url
      result['fullUrl'] = url
    else:
      result['fullUrl'] = urljoin(parentUrl, url)
      '''
      if url.startswith('/'):
        # 在根目录中
        result['fullUrl'] = pParse.scheme+'://'+pParse.netloc+url
      else:
        # 在parent的对应目录中
        if url.startswith('./'):
          url = url[2:]
        result['fullUrl'] = parentUrl[:len(parentUrl)-len(pFn)]+url
        '''
    ## 处理inRoot
    result['inRoot'] = result['fullUrl'].startswith(self._root)
    ## 处理filename和downPath
    fn = re.search(r'([^/]*)$', result['fullUrl']).group(0)
    urlPath = result['fullUrl'][:len(result['fullUrl'])-len(fn)]
    if result['inRoot']:
      # 在根目录中,按对应目录处理
      result['downPath'] = urlPath.replace(self._root, '').replace('/', '\\')
      result['filename'] = fn.replace('?', '_')
    else:
      # 不在根目录中,全部保存到_outside并处理文件名
      result['downPath'] = '\\_outside'+'\\'
      result['filename'] = urlPath.replace('http://', '')
      result['filename'] = result['filename'].replace('https://', '')
      result['filename'] = result['filename'].replace('.', '_')
      result['filename'] = result['filename'].replace('/', '_')
      result['filename'] = result['filename']+fn
    return result
   
  def _handleSoup(self, soup, url):
    """指定一个soup以及对应的url,对这个soup进行改写,
    并添加后续需要分析的链接,添加最后需要下载的资源.
    返回值是一个set(),表示下一层要分析的链接.
    这里提取把资源需要的文件夹创建了,因为如果等下载时创建,可能
    要重新通过路径分析文件夹名,有额外消耗.
    """
    links = set() # 下一层的链接
    # 指定需要处理的元素的选择器和对应的url属性
    elementSelectors = [
      ('a[href]', 'href'), ('img[src]', 'src'),
      ('link[href][rel=stylesheet][type=text\/css]', 'href'),
      ('iframe[src]', 'src'), ('[background]', 'background')
    ]
    # 有效链接的flag,如果是False则表示无效,
    # True才根据情况改写和创建文件夹
    flag = False
    # 处理这些选择,并分析其对应的url是进入下一层分析还是只需下载
    for selector, attr in elementSelectors:
      for element in soup.select(selector):
        flag = False
        parse = self._urlParse(element.attrs[attr], url)
        p = self._outputDir+parse['downPath']
        fp = p+parse['filename']
        # 对链接分析,如果是网页类则进入下一层分析,否则只加入下载表
        # 如果分析后发现不用理会,那么改写的flag就是False
        if self._urlFilter(parse['fullUrl']):
          links.add(parse['fullUrl'])
          flag = True
        else:
          if self._resourceFilter(parse['fullUrl']):
            self._resource[parse['fullUrl']] = fp
            flag = True
        # 有效链接判断,有效时新建文件夹,改写绝对链接
        if not flag:
          continue
        if (not os.path.isdir(p)):
          os.makedirs(p, exist_ok=True)
        if not parse['inRoot']:
          element.attrs[attr] = parse['downPath'][1:]+parse['filename']
        if self._root in element.attrs[attr]:
          element.attrs[attr] = element.attrs[attr].replace(self._root) 
    # 处理完成,返回下一层链接集合
    return links
     
  def _handleUrl(self, url, level, isDebug = True):
    """处理一个url,并指定其当前层数为level,若层数大于maxLevel.
    具体处理方式为先看看自己是否已经处理.如果已处理,则看看上次
    处理的层数是否更靠前,若更靠前就不用处理了,否则需要进行下载以外
    的处理并更新记录.
    
    若isDebug为True,则不下载只提示"""
    downFlag = True # 是否需要下载本页面
    ## 层数过大,跳过
    if level > self._maxLevel:
      return
    ## 已下载过,不再下载,看需不需要分析,如果不需要分析则跳过
    if url in self._htmlRecord:
      downFlag = False
      if self._htmlRecord[url] <= level:
        return
    ## 到这里则是需要分析了,是否再下载要看downFlag
    try:
      if isDebug:
        print("开始获取:"+url)
      resp = self._session.get(url=url, headers=self._headers)
      text = resp.content.decode(self._encoding)
      soup = BeautifulSoup(text, 'lxml')
    except:
      print("获取分析"+url+"失败!")
      return
    # 处理soup,改写成最后可以复制到硬盘的状态,并获取链接与资源
    links = self._handleSoup(soup, url)
    # 刷新本页面的分析层数记录,若有需要则下载本页面.下载本页面
    self._htmlRecord[url] = level
    # 保存当前页面
    if downFlag:
      hrefParse = self._urlParse(url, self._root)
      p = self._outputDir+hrefParse['downPath']
      if (not os.path.isdir(p)):
        os.makedirs(p, exist_ok=True)
      fp = p+hrefParse['filename']
      if isDebug:
        print("debug下载页面:"+url+"\n为:"+fp)
        #print(self._htmlRecord)
      else:
        with open(fp, 'wb') as f:
          f.write(soup.prettify().encode('utf-8'))
        print("下载页面成功:"+url)
    # 递归分析链接,并增加层数
    if level+1 <= self._maxLevel:
      for link in links:
        self._handleUrl(link, level+1, isDebug)
    return
    
  def _downResource(self, isDebug=True):
    """下载所有的resource,在爬虫最后使用"""
    if isDebug:
      print('debug - resource:')
      print(self._resource)
    for src in self._resource:
      fp = self._resource[src]
      if isDebug:
        print("debug下载资源:"+src+"\n为:"+fp)
      else:
        if os.path.exists(fp):
          print("资源:"+src+"已存在")
          continue
        try:
          resp = self._session.get(url=src, headers=self._headers) 
          with open(fp, 'wb') as f:
            f.write(resp.content)
          print("下载资源成功:"+src)
        except:
          print("下载资源失败:"+src)
          continue
    return
  
  def main(self, isDebug = True):
    """正式执行,如果isDebug为True,则不下载只提示"""
    os.makedirs(self._outputDir, exist_ok=True)
    print('开始执行任务!')
    for x in self._urlList:
      self._handleUrl(x[0], x[1], isDebug)
    self._downResource(isDebug)
    return

## 开启脚本    
if __name__ == '__main__':
  root = r'http://www.srworld.net/gonglue/srw/srwz/'
  urlList = [
    (r'http://www.srworld.net/gonglue/srw/srwz/srwz.htm', 1)
  ]
  outputDir = r'F:\\Download\srw_data\z'
  maxLevel = 4
  headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)',
    'Referer': urlList[0][0]
  }
  
  m = MyWebGet(root, outputDir, urlList, maxLevel, headers, 'gbk')
  isDebug = False
  m.main(isDebug)
    